package com.celan.year2023.month05.day28;

import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;

/**
 * @author Celan
 * @description TODO
 */
public class Solution02 {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int len1 = nums1.length, len2 = nums2.length;
        List<List<Integer>> res = new ArrayList<>();

        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> a[0] - b[0]);
        //nums1[0] + nums2[0] 一定为最小值
        queue.add(new int[]{nums1[0] + nums2[0], 0, 0});

        //可能成为下一个数对的是 (idx1+1,idx2) 和 (idx1,j+idx2)
        //(i,j−1) 出堆时,将 (i,j) 入堆; 而 (i−1,j) 出堆时只计入答案
        while (!queue.isEmpty() && res.size() < k) {
            int[] poll = queue.poll();
            int idx1 = poll[1], idx2 = poll[2];

            List<Integer> pair = new ArrayList<>();
            pair.add(nums1[idx1]);
            pair.add(nums2[idx2]);
            res.add(pair);

            //去重 -> 在循环的过程中去把 (idx1,0) 入堆
            if (idx2 == 0 && idx1 + 1 < len1)
                queue.add(new int[]{nums1[idx1 + 1] + nums2[0], idx1 + 1, 0});
            if (idx2 + 1 < len2)
                queue.add(new int[]{nums1[idx1] + nums2[idx2 + 1], idx1, idx2 + 1});

        }

        return res;
    }
}
